import numpy as np
import scipy.stats


def mean_confidence_interval(data, before, confidence=0.95):
    a = 1.0 * np.array(data[:before])
    n = len(a)
    m, se = np.mean(a), scipy.stats.sem(a)
    h = se * scipy.stats.t.ppf((1 + confidence) / 2.0, n - 1)
    return m, h  # m-h, m+h


def composite_k(epsilon, data_array, k):
    data_array = data_array.astype(float)
    e = np.full_like(data_array, epsilon)
    indexes = np.array(list(range(k, k + len(data_array))))
    e = np.divide(e, indexes)
    return e


def add_laplace_noise(
    data,
    sensitivity: float,
    epsilon: float,
    seed=None,
    personalized=False,
    dynamic_epsilon=False,
    composition=False,
    k=1,
) -> np.ndarray:
    """
    为数值型数据添加满足ε-差分隐私的拉普拉斯噪声

    Parameters:
        data (float/array): 原始标量值或数值型数据集
        sensitivity (float): 查询函数的敏感度，定义为相邻数据集的最大结果差异
        epsilon (float): 隐私预算，控制隐私保护强度(ε>0)

    Returns:
        noisy_data: 添加拉普拉斯噪声后的扰动数据
    """

    data_array = np.asarray(data)

    if personalized:
        # e = epsilon / len(data_array) * 0.0001 + 0.05
        alpha = 0.5
        beta = 0.5
        e = epsilon * (alpha * 0.7 + beta * 1.0 / 0.3)
        if composition:
            e = composite_k(e, data_array, k)
        scale = sensitivity / e
    else:
        if composition:
            epsilon = composite_k(epsilon, data_array, k)
        scale = sensitivity / epsilon

    if seed:
        for val in data_array:
            rng = np.random.default_rng(seed)
            noise = rng.laplace(loc=0, scale=scale, size=np.shape(val))
    else:
        # rng = np.random.default_rng(seed=123)
        # noise = rng.laplace(loc=0, scale=scale, size=np.shape(data))
        noise = np.random.laplace(loc=0.0, scale=scale, size=data_array.shape)
    return data_array + noise
